#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
代码生成器
"""

from typing import Dict, Optional

from .ir import IRModule
from .codegen.file_structure_generator import FileStructureGenerator
from .codegen.node_renderer import NodeRenderer
from .codegen.task_collector import CodeGenTaskCollector


class CodeGenerator:
    """代码生成器"""

    def __init__(self):
        """初始化代码生成器"""
        pass

    def generate_from_ir(
        self,
        ir_module: IRModule,
        header_output_dir: str = "",
        module_name: Optional[str] = None,
    ) -> Dict[str, str]:
        """从新IR生成代码"""
        if module_name is None:
            module_name = ir_module.name

        # 计算输出文件名
        header_file = f"{module_name}.h"
        source_file = f"{module_name}.cpp"

        # 先生成源文件内容（收集 includes），再生成头文件
        source_content = self._generate_source(ir_module, module_name, header_file)
        header_content = self._generate_header(ir_module, module_name)

        return {header_file: header_content, source_file: source_content}

    def _generate_header(self, ir_module: IRModule, module_name: str) -> str:
        """生成头文件内容"""
        # 收集任务
        collector = CodeGenTaskCollector()
        tasks = collector.collect_tasks(ir_module, for_source=False)

        # 创建渲染器
        render_ctx = {"type": "header"}
        node_renderer = NodeRenderer(render_ctx)

        # 生成代码
        parts = []
        for task in tasks:
            if task.task_type == "decl":
                code = node_renderer.render_node(task.node, for_source=False)
                if code:
                    parts.append(code)

        nodes_content = "\n\n".join(parts)

        # 获取 lambda 结构体（仍然需要，因为结构体是在生成时添加到 lambda_structs 的）
        lambda_structs = ir_module.get_extension_param("lambda_structs") or []

        # 生成完整的头文件
        return FileStructureGenerator.generate_header_file(
            ir_module=ir_module,
            module_name=module_name,
            nodes_content=nodes_content,
            lambda_structs=lambda_structs,
        )

    def _generate_source(
        self, ir_module: IRModule, module_name: str, header_file: str
    ) -> str:
        """生成源文件内容"""
        # 收集任务
        collector = CodeGenTaskCollector()
        tasks = collector.collect_tasks(ir_module, for_source=True)

        # 创建渲染器
        render_ctx = {"type": "source"}
        node_renderer = NodeRenderer(render_ctx)

        # 按照任务顺序直接生成代码（任务收集器已经按照正确的顺序排序）
        parts = []
        for task in tasks:
            if task.task_type == "impl":
                code = node_renderer.render_node(task.node, for_source=True)
                if code:
                    parts.append(code)

        combined_content = "\n\n".join(parts)

        # 获取 includes（源文件可能需要额外的 includes）
        includes = ir_module.get_includes()

        # 生成完整的源文件
        return FileStructureGenerator.generate_source_file(
            header_file=header_file,
            includes=includes,
            nodes_content=combined_content,
        )
